// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2008-2014 Gael Guennebaud <gael.guennebaud@inria.fr>
// Copyright (C) 2006-2008 Benoit Jacob <jacob.benoit.1@gmail.com>
// Copyright (C) 2016 Eugene Brevdo <ebrevdo@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_CWISE_TERNARY_OP_H
#define EIGEN_CWISE_TERNARY_OP_H

namespace Eigen {

namespace internal {
    template <typename TernaryOp, typename Arg1, typename Arg2, typename Arg3> struct traits<CwiseTernaryOp<TernaryOp, Arg1, Arg2, Arg3>>
    {
        // we must not inherit from traits<Arg1> since it has
        // the potential to cause problems with MSVC
        typedef typename remove_all<Arg1>::type Ancestor;
        typedef typename traits<Ancestor>::XprKind XprKind;
        enum
        {
            RowsAtCompileTime = traits<Ancestor>::RowsAtCompileTime,
            ColsAtCompileTime = traits<Ancestor>::ColsAtCompileTime,
            MaxRowsAtCompileTime = traits<Ancestor>::MaxRowsAtCompileTime,
            MaxColsAtCompileTime = traits<Ancestor>::MaxColsAtCompileTime
        };

        // even though we require Arg1, Arg2, and Arg3 to have the same scalar type
        // (see CwiseTernaryOp constructor),
        // we still want to handle the case when the result type is different.
        typedef typename result_of<TernaryOp(const typename Arg1::Scalar&, const typename Arg2::Scalar&, const typename Arg3::Scalar&)>::type Scalar;

        typedef typename internal::traits<Arg1>::StorageKind StorageKind;
        typedef typename internal::traits<Arg1>::StorageIndex StorageIndex;

        typedef typename Arg1::Nested Arg1Nested;
        typedef typename Arg2::Nested Arg2Nested;
        typedef typename Arg3::Nested Arg3Nested;
        typedef typename remove_reference<Arg1Nested>::type _Arg1Nested;
        typedef typename remove_reference<Arg2Nested>::type _Arg2Nested;
        typedef typename remove_reference<Arg3Nested>::type _Arg3Nested;
        enum
        {
            Flags = _Arg1Nested::Flags & RowMajorBit
        };
    };
}  // end namespace internal

template <typename TernaryOp, typename Arg1, typename Arg2, typename Arg3, typename StorageKind> class CwiseTernaryOpImpl;

/** \class CwiseTernaryOp
  * \ingroup Core_Module
  *
  * \brief Generic expression where a coefficient-wise ternary operator is
 * applied to two expressions
  *
  * \tparam TernaryOp template functor implementing the operator
  * \tparam Arg1Type the type of the first argument
  * \tparam Arg2Type the type of the second argument
  * \tparam Arg3Type the type of the third argument
  *
  * This class represents an expression where a coefficient-wise ternary
 * operator is applied to three expressions.
  * It is the return type of ternary operators, by which we mean only those
 * ternary operators where
  * all three arguments are Eigen expressions.
  * For example, the return type of betainc(matrix1, matrix2, matrix3) is a
 * CwiseTernaryOp.
  *
  * Most of the time, this is the only way that it is used, so you typically
 * don't have to name
  * CwiseTernaryOp types explicitly.
  *
  * \sa MatrixBase::ternaryExpr(const MatrixBase<Argument2> &, const
 * MatrixBase<Argument3> &, const CustomTernaryOp &) const, class CwiseBinaryOp,
 * class CwiseUnaryOp, class CwiseNullaryOp
  */
template <typename TernaryOp, typename Arg1Type, typename Arg2Type, typename Arg3Type>
class CwiseTernaryOp : public CwiseTernaryOpImpl<TernaryOp, Arg1Type, Arg2Type, Arg3Type, typename internal::traits<Arg1Type>::StorageKind>,
                       internal::no_assignment_operator
{
public:
    typedef typename internal::remove_all<Arg1Type>::type Arg1;
    typedef typename internal::remove_all<Arg2Type>::type Arg2;
    typedef typename internal::remove_all<Arg3Type>::type Arg3;

    typedef typename CwiseTernaryOpImpl<TernaryOp, Arg1Type, Arg2Type, Arg3Type, typename internal::traits<Arg1Type>::StorageKind>::Base Base;
    EIGEN_GENERIC_PUBLIC_INTERFACE(CwiseTernaryOp)

    typedef typename internal::ref_selector<Arg1Type>::type Arg1Nested;
    typedef typename internal::ref_selector<Arg2Type>::type Arg2Nested;
    typedef typename internal::ref_selector<Arg3Type>::type Arg3Nested;
    typedef typename internal::remove_reference<Arg1Nested>::type _Arg1Nested;
    typedef typename internal::remove_reference<Arg2Nested>::type _Arg2Nested;
    typedef typename internal::remove_reference<Arg3Nested>::type _Arg3Nested;

    EIGEN_DEVICE_FUNC
    EIGEN_STRONG_INLINE CwiseTernaryOp(const Arg1& a1, const Arg2& a2, const Arg3& a3, const TernaryOp& func = TernaryOp())
        : m_arg1(a1), m_arg2(a2), m_arg3(a3), m_functor(func)
    {
        // require the sizes to match
        EIGEN_STATIC_ASSERT_SAME_MATRIX_SIZE(Arg1, Arg2)
        EIGEN_STATIC_ASSERT_SAME_MATRIX_SIZE(Arg1, Arg3)

        // The index types should match
        EIGEN_STATIC_ASSERT((internal::is_same<typename internal::traits<Arg1Type>::StorageKind, typename internal::traits<Arg2Type>::StorageKind>::value),
                            STORAGE_KIND_MUST_MATCH)
        EIGEN_STATIC_ASSERT((internal::is_same<typename internal::traits<Arg1Type>::StorageKind, typename internal::traits<Arg3Type>::StorageKind>::value),
                            STORAGE_KIND_MUST_MATCH)

        eigen_assert(a1.rows() == a2.rows() && a1.cols() == a2.cols() && a1.rows() == a3.rows() && a1.cols() == a3.cols());
    }

    EIGEN_DEVICE_FUNC
    EIGEN_STRONG_INLINE Index rows() const
    {
        // return the fixed size type if available to enable compile time
        // optimizations
        if (internal::traits<typename internal::remove_all<Arg1Nested>::type>::RowsAtCompileTime == Dynamic &&
            internal::traits<typename internal::remove_all<Arg2Nested>::type>::RowsAtCompileTime == Dynamic)
            return m_arg3.rows();
        else if (internal::traits<typename internal::remove_all<Arg1Nested>::type>::RowsAtCompileTime == Dynamic &&
                 internal::traits<typename internal::remove_all<Arg3Nested>::type>::RowsAtCompileTime == Dynamic)
            return m_arg2.rows();
        else
            return m_arg1.rows();
    }
    EIGEN_DEVICE_FUNC
    EIGEN_STRONG_INLINE Index cols() const
    {
        // return the fixed size type if available to enable compile time
        // optimizations
        if (internal::traits<typename internal::remove_all<Arg1Nested>::type>::ColsAtCompileTime == Dynamic &&
            internal::traits<typename internal::remove_all<Arg2Nested>::type>::ColsAtCompileTime == Dynamic)
            return m_arg3.cols();
        else if (internal::traits<typename internal::remove_all<Arg1Nested>::type>::ColsAtCompileTime == Dynamic &&
                 internal::traits<typename internal::remove_all<Arg3Nested>::type>::ColsAtCompileTime == Dynamic)
            return m_arg2.cols();
        else
            return m_arg1.cols();
    }

    /** \returns the first argument nested expression */
    EIGEN_DEVICE_FUNC
    const _Arg1Nested& arg1() const { return m_arg1; }
    /** \returns the first argument nested expression */
    EIGEN_DEVICE_FUNC
    const _Arg2Nested& arg2() const { return m_arg2; }
    /** \returns the third argument nested expression */
    EIGEN_DEVICE_FUNC
    const _Arg3Nested& arg3() const { return m_arg3; }
    /** \returns the functor representing the ternary operation */
    EIGEN_DEVICE_FUNC
    const TernaryOp& functor() const { return m_functor; }

protected:
    Arg1Nested m_arg1;
    Arg2Nested m_arg2;
    Arg3Nested m_arg3;
    const TernaryOp m_functor;
};

// Generic API dispatcher
template <typename TernaryOp, typename Arg1, typename Arg2, typename Arg3, typename StorageKind>
class CwiseTernaryOpImpl : public internal::generic_xpr_base<CwiseTernaryOp<TernaryOp, Arg1, Arg2, Arg3>>::type
{
public:
    typedef typename internal::generic_xpr_base<CwiseTernaryOp<TernaryOp, Arg1, Arg2, Arg3>>::type Base;
};

}  // end namespace Eigen

#endif  // EIGEN_CWISE_TERNARY_OP_H
